import torch
import torch.optim as optim
from models import Actor, Critic

class MADDPGAgent:
    def __init__(self, agent_id, obs_dim, action_dim, total_obs_dim, total_action_dim,
                 actor_lr=1e-3, critic_lr=1e-3, gamma=0.95, tau=0.01):
        self.agent_id = agent_id
        self.obs_dim = obs_dim
        self.action_dim = action_dim
        self.total_obs_dim = total_obs_dim
        self.total_action_dim = total_action_dim
        self.gamma = gamma
        self.tau = tau

        self.actor = Actor(obs_dim, action_dim)
        self.target_actor = Actor(obs_dim, action_dim)
        self.critic = Critic(total_obs_dim, total_action_dim)
        self.target_critic = Critic(total_obs_dim, total_action_dim)

        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=actor_lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=critic_lr)

        self.target_actor.load_state_dict(self.actor.state_dict())
        self.target_critic.load_state_dict(self.critic.state_dict())

    def update_target(self):
        for target_param, param in zip(self.target_actor.parameters(), self.actor.parameters()):
            target_param.data.copy_(self.tau * param.data + (1.0 - self.tau) * target_param.data)
        for target_param, param in zip(self.target_critic.parameters(), self.critic.parameters()):
            target_param.data.copy_(self.tau * param.data + (1.0 - self.tau) * target_param.data)
